#!/usr/bin/env python
######################################################################
##
## Copyright (C) 2008,  Simon Kagstrom
##
## Filename:      shcov
## Author:        Simon Kagstrom <simon.kagstrom@gmail.com>
## Description:   Test shell script coverage
##
## $Id:$
##
######################################################################

# export PS4='SHCOV:::${BASH_SOURCE}:::${LINENO}::: SHCOV:'

import os, sys, subprocess, pickle, signal, select

base_dir = os.path.abspath(sys.path[0] + "/../")
sys.path =  [base_dir] + sys.path

import shcov

from shcov.file import File

class ShcovCollector:
    def __init__(self, args, outpath = "/tmp/shcov", shell = ["bash", "-x"]):
	self.files = {}
	self.args = args
	self.outpath = outpath
	self.shell = shell

	# Where we execute this script
	os.environ["PS4"] = "SHCOV:::${BASH_SOURCE}:::${LINENO}::: SHCOV:"

	self.process = None

    def handle_line(self, line):
        """Handle one line of stderr input"""
        if line.find("SHCOV:::") == -1:
            sys.stderr.write(line)
            return
	parts = line.split(":::")

	file_name = parts[1]
	line_nr = parts[2]

	# Get the file for this path
	path = os.path.abspath(file_name)
	try:
	    file = self.files[path]
	except KeyError, e:
	    # Not found, try to load it or otherwise create a new one
	    try:
		file = shcov.file.load(self.outpath + path + ".pkl")
	    except Exception, e:
		# Create a new one (no such file)
                try:
                    file = File(path)
                except IOError, e:
                    # Cannot create the object (file not found most likely),
                    # just ignore it
                    return
            # And create the dir to save it in
	    try:
		os.makedirs( os.path.dirname(self.outpath + file.path) )
	    except OSError, e:
		pass
	    self.files[path] = file
	file.add_to_line(line_nr)
        file.save(self.outpath + file.path + ".pkl")

    def parse_output(self):
        """Process the bash stderr output"""
	self.process = subprocess.Popen( self.shell + self.args,  stdin = sys.stdin,
					 stdout = sys.stdout, stderr = subprocess.PIPE)

        fd_poll = select.poll()
        fd_poll.register( self.process.stderr.fileno() )

	while True:
            # Did the process exit?
            ret = self.process.poll()
            if ret != None:
                try:
                    pid = os.fork()
                except OSError, e:
                    raise Exception, "%s [%d]" % (e.strerror, e.errno)
                if pid == 0:
                    # Read the last stuff
                    for line in self.process.stderr.readlines():
                        self.handle_line(line)
                    # Save everything
                    self.save()
                return ret

	    # Write out everything on stderr
	    self.process.stderr.flush()

            # If stderr is not ready, just continue
            status = fd_poll.poll(10)
            if status == []:
                continue

	    line = self.process.stderr.readline()
            self.handle_line(line)

    def save(self):
	for file in self.files.values():
	    file.save(self.outpath + file.path + ".pkl")

def usage():
    print "Usage: shcov [-h] [--output=where] [--shell=what] script...\n"
    print "Produce coverage data for 'script'. Options are\n"
    print "  --output=where  write data to 'where' instead of /tmp/shcov"
    print "  --shell=what    ues 'what' (including arguments) as shell instead of 'bash -x'"

    sys.exit(1)

glob_sc = None

def sighandler(signum, frame):
    glob_sc.save()
    sys.exit(1)

if __name__ == "__main__":
    if len(sys.argv) < 2:
	usage()

    outpath = "/tmp/shcov"
    shell = ["bash", "-x"]

    # Really, really, ugly, but it's done this way to avoid parsing the
    # options passed to the shell script
    last = 1
    for arg in sys.argv[1:min(len(sys.argv), 3)]:
        if arg == "-h" or arg == "--help":
            usage()
        elif arg.startswith("--output="):
            last = last + 1
            outpath = arg.strip("--output=")
        elif arg.startswith("--shell="):
            last = last + 1
            shell = arg.strip("--shell=").split()
        else:
            # First non-option
            break
    args = sys.argv[last:]

    if len(args) < 1:
	usage()

    sc = ShcovCollector(args, outpath = outpath, shell = shell)
    glob_sc = sc
    signal.signal(signal.SIGTERM, sighandler)
    out = 1
    try:
	out = sc.parse_output()
    except KeyboardInterrupt, e:
	# Just save
        sc.save()
    sys.exit(out)
